package org.yagnus.ml.ie;

import java.util.ArrayList;
import java.util.HashSet;

import org.yagnus.datastructure.Range;

/**
 * This class implements PR calculation based on ranges. The algorithm is O(T*E) where T is the
 * number of ground truth segments and E is the number of extracted segments
 * 
 * Allows for two modes of operation: retrieval means exact equivalence or intersection
 * 
 * @author hc.busy
 * 
 * @param <B>
 *            The base number type
 * @param <R>
 *            The range type that we perform PR over
 */
public class PrecisionRecallCalculator<B extends Comparable<B>, R extends Range<R, B>> extends PrecisionRecall {

    ArrayList<R> truths = new ArrayList<R>();
    ArrayList<R> extractions = new ArrayList<R>();
    HashSet<Integer> recalled = new HashSet<Integer>();
    HashSet<Integer> truePositives = new HashSet<Integer>();

    final boolean exact;

    public PrecisionRecallCalculator() {
        super();
        reset();
        exact = false;
    }

    public PrecisionRecallCalculator(boolean exact) {
        super();
        reset();
        this.exact = exact;
    }

    public boolean reset() {
        truths.clear();
        extractions.clear();
        recalled.clear();
        truePositives.clear();
        return true;
    }

    /**
     * @param the
     *            instance of groundTruth to add
     */
    public void addGroundTruth(R t) {
        boolean recalled = false;
        for (int i = 0; i < extractions.size(); ++i) {
            if (rangeRecalled(t, extractions.get(i))) {
                recalled = true;
                truePositives.add(i);
            }

        }

        truths.add(t);

        if (recalled) {
            this.recalled.add(truths.size() - 1);
        }

    }

    /**
     * @param e
     *            the instance of extraction to add
     */
    public void addExtraction(R e) {
        boolean truePositive = false;

        for (int i = 0; i < truths.size(); ++i) {
            if (rangeRecalled(truths.get(i), e)) {
                truePositive = true;
                recalled.add(i);
            }
        }

        extractions.add(e);
        if (truePositive) {
            truePositives.add(extractions.size() - 1);
        }
    }

    public boolean rangeRecalled(R truth, R extraction) {
        if (extraction == null || truth == null)
            return false;
        if (exact) {
            return truth.equals(extraction);
        } else {
            return truth.intersects(extraction);
        }
    }

    @Override
    public double getTruePositive() {
        return truePositives.size();
    }

    @Override
    public double getFalsePositive() {
        return extractions.size() - truePositives.size();
    }

    // True negative is always zero in this calculator since there is no way to
    // add one
    @Override
    public double getTrueNegative() {
        return 0;
    }

    @Override
    public double getFalseNegative() {
        return truths.size() - recalled.size();
    }

    @Override
    public double getRecalled() {
        return recalled.size();
    }

    @Override
    public double getTruthCount() {
        return truths.size();
    }

    @Override
    public double getExtractionCount() {
        return extractions.size();
    }

    public boolean isExact() {
        return exact;
    }

}
